#!/usr/bin/env python3
"""
边缘保留的灰度转换方法

针对边缘检测任务优化的灰度转换方法，能够更好地保留边缘信息
"""

import cv2
import numpy as np


def rgb_to_gray_edge_preserving(image_bgr, method='max_gradient'):
    """
    将BGR图像转换为灰度图，同时保留边缘信息
    
    Args:
        image_bgr: BGR格式的图像 (H, W, 3)
        method: 转换方法
            - 'luminance': 标准亮度加权 Y = 0.114*B + 0.587*G + 0.299*R (默认cv2.COLOR_BGR2GRAY)
            - 'max_gradient': 保留三通道中梯度最大的区域（最优边缘保留）
            - 'gradient_fusion': 融合三通道的梯度信息
            - 'max_channel': 每个像素取三通道的最大值（高对比度）
            - 'desaturation': 取三通道的平均值（简单快速）
    
    Returns:
        gray: 灰度图像 (H, W, 1)
    """
    
    if method == 'luminance':
        # 标准亮度加权（OpenCV默认）
        gray = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY)
        
    elif method == 'max_gradient':
        # 方法1：保留梯度最大的通道区域
        # 这是最适合边缘检测的方法
        b, g, r = cv2.split(image_bgr)
        
        # 计算每个通道的梯度幅值
        grad_b = np.abs(cv2.Sobel(b, cv2.CV_32F, 1, 0, ksize=3)) + \
                 np.abs(cv2.Sobel(b, cv2.CV_32F, 0, 1, ksize=3))
        grad_g = np.abs(cv2.Sobel(g, cv2.CV_32F, 1, 0, ksize=3)) + \
                 np.abs(cv2.Sobel(g, cv2.CV_32F, 0, 1, ksize=3))
        grad_r = np.abs(cv2.Sobel(r, cv2.CV_32F, 1, 0, ksize=3)) + \
                 np.abs(cv2.Sobel(r, cv2.CV_32F, 0, 1, ksize=3))
        
        # 创建梯度掩码：每个像素位置选择梯度最大的通道
        grad_stack = np.stack([grad_b, grad_g, grad_r], axis=-1)
        max_grad_idx = np.argmax(grad_stack, axis=-1)
        
        # 根据梯度掩码选择对应通道的像素值
        gray = np.zeros_like(b, dtype=np.float32)
        gray[max_grad_idx == 0] = b[max_grad_idx == 0]
        gray[max_grad_idx == 1] = g[max_grad_idx == 1]
        gray[max_grad_idx == 2] = r[max_grad_idx == 2]
        
        gray = gray.astype(np.uint8)
        
    elif method == 'gradient_fusion':
        # 方法2：融合三通道的梯度信息
        # 计算每个通道的梯度，然后加权融合
        b, g, r = cv2.split(image_bgr)
        
        # 计算梯度
        grad_b_x = cv2.Sobel(b, cv2.CV_32F, 1, 0, ksize=3)
        grad_b_y = cv2.Sobel(b, cv2.CV_32F, 0, 1, ksize=3)
        grad_g_x = cv2.Sobel(g, cv2.CV_32F, 1, 0, ksize=3)
        grad_g_y = cv2.Sobel(g, cv2.CV_32F, 0, 1, ksize=3)
        grad_r_x = cv2.Sobel(r, cv2.CV_32F, 1, 0, ksize=3)
        grad_r_y = cv2.Sobel(r, cv2.CV_32F, 0, 1, ksize=3)
        
        # 融合梯度（使用亮度权重）
        grad_x = 0.114 * grad_b_x + 0.587 * grad_g_x + 0.299 * grad_r_x
        grad_y = 0.114 * grad_b_y + 0.587 * grad_g_y + 0.299 * grad_r_y
        
        # 从融合梯度重建图像
        # 使用标准亮度转换作为基础，融合梯度信息增强边缘
        gray_base = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2GRAY).astype(np.float32)
        grad_magnitude = np.sqrt(grad_x**2 + grad_y**2)
        
        # 归一化梯度幅值
        if grad_magnitude.max() > 0:
            grad_magnitude = grad_magnitude / grad_magnitude.max() * 50  # 增强系数
        
        # 叠加梯度到基础灰度图
        gray = gray_base + grad_magnitude
        gray = np.clip(gray, 0, 255).astype(np.uint8)
        
    elif method == 'max_channel':
        # 方法3：每个像素取三通道的最大值
        # 这会产生高对比度，有助于边缘检测
        b, g, r = cv2.split(image_bgr)
        gray = np.maximum(np.maximum(b, g), r)
        
    elif method == 'desaturation':
        # 方法4：简单的通道平均（去饱和度）
        gray = np.mean(image_bgr, axis=2).astype(np.uint8)
        
    else:
        raise ValueError(f"未知的转换方法: {method}")
    
    # 确保输出是 (H, W, 1) 形状
    if len(gray.shape) == 2:
        gray = np.expand_dims(gray, axis=-1)
    
    return gray


def compare_grayscale_methods(image_bgr, save_path=None):
    """
    比较不同灰度转换方法的效果
    
    Args:
        image_bgr: BGR格式的输入图像
        save_path: 如果提供，将比较结果保存为图像
    
    Returns:
        results: 字典，包含不同方法的转换结果
    """
    methods = ['luminance', 'max_gradient', 'gradient_fusion', 'max_channel', 'desaturation']
    results = {}
    
    for method in methods:
        results[method] = rgb_to_gray_edge_preserving(image_bgr, method)
    
    if save_path:
        import matplotlib.pyplot as plt
        
        fig, axes = plt.subplots(2, 3, figsize=(15, 10))
        axes = axes.ravel()
        
        # 显示原图
        axes[0].imshow(cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB))
        axes[0].set_title('Original (BGR)')
        axes[0].axis('off')
        
        # 显示各种方法
        for idx, method in enumerate(methods, start=1):
            axes[idx].imshow(results[method].squeeze(), cmap='gray')
            axes[idx].set_title(method)
            axes[idx].axis('off')
        
        plt.tight_layout()
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        print(f"比较结果已保存至: {save_path}")
    
    return results


if __name__ == '__main__':
    """
    测试脚本：比较不同灰度转换方法的效果
    """
    import sys
    
    if len(sys.argv) < 2:
        print("用法: python edge_preserving_gray.py <image_path> [output_path]")
        print("示例: python edge_preserving_gray.py test.jpg comparison.png")
        sys.exit(1)
    
    image_path = sys.argv[1]
    output_path = sys.argv[2] if len(sys.argv) > 2 else 'grayscale_comparison.png'
    
    # 读取图像
    image = cv2.imread(image_path, cv2.IMREAD_COLOR)
    if image is None:
        print(f"错误：无法读取图像 {image_path}")
        sys.exit(1)
    
    print(f"图像尺寸: {image.shape}")
    print(f"正在比较不同灰度转换方法...")
    
    # 比较方法
    results = compare_grayscale_methods(image, save_path=output_path)
    
    print("\n各方法特点：")
    print("1. luminance: 标准亮度加权（OpenCV默认），平衡性好")
    print("2. max_gradient: 保留梯度最大的通道（★推荐用于边缘检测）")
    print("3. gradient_fusion: 融合三通道梯度，增强边缘")
    print("4. max_channel: 高对比度，边缘明显但可能有噪声")
    print("5. desaturation: 简单平均，速度快但边缘保留一般")
